/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package org.apache.iotdb.db.mpp.execution.operator.source;

import org.apache.iotdb.db.engine.querycontext.QueryDataSource;
import org.apache.iotdb.db.mpp.aggregation.Aggregator;
import org.apache.iotdb.db.mpp.aggregation.timerangeiterator.ITimeRangeIterator;
import org.apache.iotdb.db.mpp.execution.operator.OperatorContext;
import org.apache.iotdb.db.mpp.plan.planner.plan.node.PlanNodeId;
import org.apache.iotdb.db.mpp.plan.planner.plan.parameter.GroupByTimeParameter;
import org.apache.iotdb.tsfile.common.conf.TSFileDescriptor;
import org.apache.iotdb.tsfile.file.metadata.enums.TSDataType;
import org.apache.iotdb.tsfile.file.metadata.statistics.Statistics;
import org.apache.iotdb.tsfile.read.common.TimeRange;
import org.apache.iotdb.tsfile.read.common.block.TsBlock;
import org.apache.iotdb.tsfile.read.common.block.TsBlockBuilder;
import org.apache.iotdb.tsfile.utils.Pair;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.TimeUnit;

import static org.apache.iotdb.db.mpp.execution.operator.AggregationUtil.*;

public abstract class AbstractSeriesAggregationScanOperator implements DataSourceOperator {

    protected final PlanNodeId sourceId;
    protected final OperatorContext operatorContext;
    protected final boolean ascending;
    protected final boolean isGroupByQuery;

    protected SeriesScanUtil seriesScanUtil;
    protected int subSensorSize;

    protected TsBlock inputTsBlock;

    protected final ITimeRangeIterator timeRangeIterator;
    // current interval of aggregation window [curStartTime, curEndTime)
    protected TimeRange curTimeRange;

    // We still think aggregator in SeriesAggregateScanOperator is a inputRaw step.
    // But in facing of statistics, it will invoke another method processStatistics()
    protected final List<Aggregator> aggregators;

    // using for building result tsBlock
    protected final TsBlockBuilder resultTsBlockBuilder;

    protected boolean finished = false;

    private final long maxRetainedSize;
    private final long maxReturnSize;

    public AbstractSeriesAggregationScanOperator(
            PlanNodeId sourceId,
            OperatorContext context,
            SeriesScanUtil seriesScanUtil,
            int subSensorSize,
            List<Aggregator> aggregators,
            ITimeRangeIterator timeRangeIterator,
            boolean ascending,
            GroupByTimeParameter groupByTimeParameter,
            long maxReturnSize) {
        this.sourceId = sourceId;
        this.operatorContext = context;
        this.ascending = ascending;
        this.isGroupByQuery = groupByTimeParameter != null;
        this.seriesScanUtil = seriesScanUtil;
        this.subSensorSize = subSensorSize;
        this.aggregators = aggregators;
        this.timeRangeIterator = timeRangeIterator;

        List<TSDataType> dataTypes = new ArrayList<>();
        for (Aggregator aggregator : aggregators) {
            dataTypes.addAll(Arrays.asList(aggregator.getOutputType()));
        }
        this.resultTsBlockBuilder = new TsBlockBuilder(dataTypes);

        this.maxRetainedSize =
                (1L + subSensorSize) * TSFileDescriptor.getInstance().getConfig().getPageSizeInByte();
        this.maxReturnSize = maxReturnSize;
    }

    @Override
    public PlanNodeId getSourceId() {
        return sourceId;
    }

    @Override
    public OperatorContext getOperatorContext() {
        return operatorContext;
    }

    @Override
    public void initQueryDataSource(QueryDataSource dataSource) {
        seriesScanUtil.initQueryDataSource(dataSource);
    }

    @Override
    public long calculateMaxPeekMemory() {
        return maxRetainedSize + maxReturnSize;
    }

    @Override
    public long calculateMaxReturnSize() {
        return maxReturnSize;
    }

    @Override
    public long calculateRetainedSizeAfterCallingNext() {
        return maxRetainedSize;
    }

    @Override
    public boolean hasNext() {
        return timeRangeIterator.hasNextTimeRange();
    }

    @Override
    public TsBlock next() {
        // start stopwatch
        long maxRuntime = operatorContext.getMaxRunTime().roundTo(TimeUnit.NANOSECONDS);
        long start = System.nanoTime();

        while (System.nanoTime() - start < maxRuntime
                && timeRangeIterator.hasNextTimeRange()
                && !resultTsBlockBuilder.isFull()) {
            // move to next time window
            curTimeRange = timeRangeIterator.nextTimeRange();

            // clear previous aggregation result
            for (Aggregator aggregator : aggregators) {
                aggregator.updateTimeRange(curTimeRange);
            }

            // calculate aggregation result on current time window
            calculateNextAggregationResult();
        }

        if (resultTsBlockBuilder.getPositionCount() > 0) {
            TsBlock resultTsBlock = resultTsBlockBuilder.build();
            resultTsBlockBuilder.reset();
            return resultTsBlock;
        } else {
            return null;
        }
    }

    @Override
    public boolean isFinished() {
        return finished || (finished = !hasNextWithTimer());
    }

    protected void calculateNextAggregationResult() {
        try {
            if (calcFromCachedData()) {
                updateResultTsBlock();
                return;
            }

            // read page data firstly
            if (readAndCalcFromPage()) {
                updateResultTsBlock();
                return;
            }

            // read chunk data secondly
            if (readAndCalcFromChunk()) {
                updateResultTsBlock();
                return;
            }

            // read from file
            if (readAndCalcFromFile()) {
                updateResultTsBlock();
                return;
            }

            updateResultTsBlock();
        } catch (IOException e) {
            throw new RuntimeException("Error while scanning the file", e);
        }
    }

    protected void updateResultTsBlock() {
        appendAggregationResult(
                resultTsBlockBuilder, aggregators, timeRangeIterator.currentOutputTime());
    }

    protected boolean calcFromCachedData() {
        return calcFromRawData(inputTsBlock);
    }

    private boolean calcFromRawData(TsBlock tsBlock) {
        Pair<Boolean, TsBlock> calcResult =
                calculateAggregationFromRawData(tsBlock, aggregators, curTimeRange, ascending);
        inputTsBlock = calcResult.getRight();
        return calcResult.getLeft();
    }

    protected void calcFromStatistics(Statistics[] statistics) {
        for (Aggregator aggregator : aggregators) {
            if (aggregator.hasFinalResult()) {
                continue;
            }
            aggregator.processStatistics(statistics);
        }
    }

    protected boolean readAndCalcFromFile() throws IOException {
        while (seriesScanUtil.hasNextFile()) {
            if (canUseCurrentFileStatistics()) {
                Statistics fileTimeStatistics = seriesScanUtil.currentFileTimeStatistics();
                if (fileTimeStatistics.getStartTime() > curTimeRange.getMax()) {
                    if (ascending) {
                        return true;
                    } else {
                        seriesScanUtil.skipCurrentFile();
                        continue;
                    }
                }
                // calc from fileMetaData
                if (curTimeRange.contains(
                        fileTimeStatistics.getStartTime(), fileTimeStatistics.getEndTime())) {
                    Statistics[] statisticsList = new Statistics[subSensorSize];
                    for (int i = 0; i < subSensorSize; i++) {
                        statisticsList[i] = seriesScanUtil.currentFileStatistics(i);
                    }
                    calcFromStatistics(statisticsList);
                    seriesScanUtil.skipCurrentFile();
                    if (isAllAggregatorsHasFinalResult(aggregators) && !isGroupByQuery) {
                        return true;
                    } else {
                        continue;
                    }
                }
            }

            // read chunk
            if (readAndCalcFromChunk()) {
                return true;
            }
        }

        return false;
    }

    protected boolean readAndCalcFromChunk() throws IOException {
        while (seriesScanUtil.hasNextChunk()) {
            if (canUseCurrentChunkStatistics()) {
                Statistics chunkTimeStatistics = seriesScanUtil.currentChunkTimeStatistics();
                if (chunkTimeStatistics.getStartTime() > curTimeRange.getMax()) {
                    if (ascending) {
                        return true;
                    } else {
                        seriesScanUtil.skipCurrentChunk();
                        continue;
                    }
                }
                // calc from chunkMetaData
                if (curTimeRange.contains(
                        chunkTimeStatistics.getStartTime(), chunkTimeStatistics.getEndTime())) {
                    // calc from chunkMetaData
                    Statistics[] statisticsList = new Statistics[subSensorSize];
                    for (int i = 0; i < subSensorSize; i++) {
                        statisticsList[i] = seriesScanUtil.currentChunkStatistics(i);
                    }
                    calcFromStatistics(statisticsList);
                    seriesScanUtil.skipCurrentChunk();
                    if (isAllAggregatorsHasFinalResult(aggregators) && !isGroupByQuery) {
                        return true;
                    } else {
                        continue;
                    }
                }
            }

            // read page
            if (readAndCalcFromPage()) {
                return true;
            }
        }
        return false;
    }

    protected boolean readAndCalcFromPage() throws IOException {
        while (seriesScanUtil.hasNextPage()) {
            if (canUseCurrentPageStatistics()) {
                Statistics pageTimeStatistics = seriesScanUtil.currentPageTimeStatistics();
                // There is no more eligible points in current time range
                if (pageTimeStatistics.getStartTime() > curTimeRange.getMax()) {
                    if (ascending) {
                        return true;
                    } else {
                        seriesScanUtil.skipCurrentPage();
                        continue;
                    }
                }
                // can use pageHeader
                if (curTimeRange.contains(
                        pageTimeStatistics.getStartTime(), pageTimeStatistics.getEndTime())) {
                    Statistics[] statisticsList = new Statistics[subSensorSize];
                    for (int i = 0; i < subSensorSize; i++) {
                        statisticsList[i] = seriesScanUtil.currentPageStatistics(i);
                    }
                    calcFromStatistics(statisticsList);
                    seriesScanUtil.skipCurrentPage();
                    if (isAllAggregatorsHasFinalResult(aggregators) && !isGroupByQuery) {
                        return true;
                    } else {
                        continue;
                    }
                }
            }

            // calc from page data
            TsBlock tsBlock = seriesScanUtil.nextPage();
            if (tsBlock == null || tsBlock.isEmpty()) {
                continue;
            }

            // calc from raw data
            if (calcFromRawData(tsBlock)) {
                return true;
            }
        }
        return false;
    }

    protected boolean canUseCurrentFileStatistics() throws IOException {
        Statistics fileStatistics = seriesScanUtil.currentFileTimeStatistics();
        return !seriesScanUtil.isFileOverlapped()
                && fileStatistics.containedByTimeFilter(seriesScanUtil.getTimeFilter())
                && !seriesScanUtil.currentFileModified();
    }

    protected boolean canUseCurrentChunkStatistics() throws IOException {
        Statistics chunkStatistics = seriesScanUtil.currentChunkTimeStatistics();
        return !seriesScanUtil.isChunkOverlapped()
                && chunkStatistics.containedByTimeFilter(seriesScanUtil.getTimeFilter())
                && !seriesScanUtil.currentChunkModified();
    }

    protected boolean canUseCurrentPageStatistics() throws IOException {
        Statistics currentPageStatistics = seriesScanUtil.currentPageTimeStatistics();
        if (currentPageStatistics == null) {
            return false;
        }
        return !seriesScanUtil.isPageOverlapped()
                && currentPageStatistics.containedByTimeFilter(seriesScanUtil.getTimeFilter())
                && !seriesScanUtil.currentPageModified();
    }
}
